from dataPre import dataPre
from configure import *
from easySample import easySample
import numpy as np
import torch
import shelve
from sklearn.metrics import roc_auc_score
import random
from siamese_img_model import imgSiamese
import sys
import os
from datetime import datetime
sys.path.append(r"/home/cyw/projects/function_sim_project/basic_script")


#   训练流程
#   我去看看原论文，这里改成使用原论文的超参数
epochs = 2000  # 高了就会过拟合了,跑了300轮，就已经明显比最优的值低不少了
BATCH_SIZE = 6     # 每个Batch使用的图像数量
# BATCH_SIZE = 1     # 每个Batch使用的图像数量、
lr = 0.00006         # 优化算法的学习率
weightDecay = 0.00001
imgPath = "/home/cyw/projects/function_sim_project/all_data/sampleDatas/img"
# pairInfs="/home/cyw/projects/function_sim_project/all_data/pair_infs/pair_infs"
pairInfs = "/home/cyw/projects/function_sim_project/all_data/EECG_pair_infs/pair_infs"
modelName = "siamese_img"


def minmaxscaler(data):
    """
        归一化
        0--255
    """
    return (data - 0)/(255)

# def generate_best_model_auc():
#     """
#         用于生成最佳模型的roc曲线
#         感觉没有调用过
#     """
#     best_path="/home/chenyongwei/project/malSimModels/model/img/LeNet_lr_6e-05_auc_0.8104091735313456.pth"
#     model=imgSiamese()
#     model.load_state_dict(torch.load(best_path))
#     model.eval()
#     pairDatas=get_data()
#     validData=pairDatas[1]
#     auc=test_auc(validData,model,saveData="True")
#     print("auc:{},数据已保存".format(auc))


class siamese_img_train():
    def __init__(self) -> None:
        self.easy_sample = easySample()

    def get_pair_infs(self, name):
        datapre = dataPre()
        data_pair_infs = datapre.get_pair_infs()
        if name == "train":
            return data_pair_infs["train"]
        elif name == "test":
            return data_pair_infs["test"]
        elif name == "valid":
            return data_pair_infs["valid"]

    def get_data(self, pairs):
        """
            获取数据
        """
        datas = [[], [], []]
        # 获得样本
        for pair in pairs:
            x = self.easy_sample.get_sample(pair[0], modelName)
            y = self.easy_sample.get_sample(pair[1], modelName)
            sim = pair[2]
            datas[0].append(x)
            datas[1].append(y)
            datas[2].append(sim)
        return [torch.tensor(datas[0]), torch.tensor(datas[1]), torch.tensor(datas[2])]

    def trans_pair_into_batch(self, datas):
        """
            转batch，调整大小到一致
        """
        res = []
        lth = len(datas[0])
        print(lth)
        a = (lth//BATCH_SIZE)+1
        for j in range(a):
            b_x, b_y, b_s = [], [], []
            for k in range(BATCH_SIZE):
                ind = BATCH_SIZE*j+k
                if ind < lth:
                    b_x.append(datas[0][ind])
                    b_y.append(datas[1][ind])
                    b_s.append(datas[2][ind])
            if len(b_x) == 0:  # 避免空数组的计算
                continue
            b_x = torch.stack((b_x))
            b_x = b_x.unsqueeze(1)
            b_y = torch.stack((b_y))
            b_y = b_y.unsqueeze(1)
            b_s = torch.stack((b_s))

            # 添加归一化操作
            b_x = minmaxscaler(np.array(b_x))
            b_x = torch.from_numpy(b_x)
            b_y = minmaxscaler(np.array(b_y))
            b_y = torch.from_numpy(b_y)
            b_s = torch.from_numpy(np.array(b_s))
            res.append((b_x, b_y, b_s))
        return res

    def train(self):
        model = imgSiamese()
        print("loading data")
        pairs = self.get_pair_infs("train")
        datas = self.get_data(pairs)
        pairs1 = self.get_pair_infs("test")
        datas1 = self.get_data(pairs1)
        pairs2 = self.get_pair_infs("valid")
        datas2 = self.get_data(pairs2)
        print("trans data into batch")
        best_auc = 0.5
        random.seed(666)
        trainData = self.trans_pair_into_batch(datas)
        testData = self.trans_pair_into_batch(datas1)
        validData = self.trans_pair_into_batch(datas2)
        self.write_to_log_file(siamese_img_train_logs,
                               "-------------start------------")
        self.write_to_log_file(
            siamese_img_train_logs, "epochs:{}  batch_size:{} lr:{}".format(epochs, BATCH_SIZE, lr))
        a = self.test_auc(validData, model)
        self.write_to_log_file(siamese_img_train_logs,
                               "initial valid  auc:{}".format(a))
        a = self.test_auc(trainData, model)
        self.write_to_log_file(siamese_img_train_logs,
                               "initial train  auc:{}".format(a))

        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        model.to(device)
        # 这个weight_decay相当于l2正则化
        optimizer = torch.optim.Adam(
            model.parameters(), lr=lr, weight_decay=weightDecay)

        epoch = 0
        while epoch < epochs:
            self.write_to_log_file(
                siamese_img_train_logs, "EPOCHS:{}".format(epoch))
            random.shuffle(trainData)
            model.train()
            running_loss = 0.0
            for step, data in enumerate(trainData):
                x, y, sim = data
                x = x.to(device, torch.float)
                y = y.to(device, torch.float)
                sim = sim.to(device, torch.float)

                pred = model(x, y)  # pred也全是0
                mse_loss = torch.nn.functional.mse_loss(pred, sim)

                optimizer.zero_grad()
                mse_loss.backward()
                optimizer.step()

                running_loss += mse_loss
            # 训练完一个epochs
            epoch += 1
            valid_a = self.test_auc(validData, model)
            self.write_to_log_file(
                siamese_img_train_logs, "     valid auc:{}    loss:{}".format(valid_a, running_loss))
            if valid_a > best_auc:
                best_auc = valid_a
                # 保存数据
                valid_a = self.test_auc(validData, model, True)
                torch.save(
                    model, r'/home/cyw/projects/function_sim_project/all_data/models/{}_model.pth'.format(modelName))

    def write_to_log_file(self, file_name_path, log_str, splitSymbol="   ", print_flag=True):
        if print_flag:
            print(log_str)
        if log_str is None:
            log_str = 'None'
        if os.path.isfile(file_name_path):
            with open(file_name_path, 'a+') as log_file:
                log_file.write(str(datetime.now())+splitSymbol+log_str + '\n')
        else:
            with open(file_name_path, 'w+') as log_file:
                log_file.write(str(datetime.now())+splitSymbol+log_str + '\n')

    def auc(self, y_truth, y_pred):
        auc_score = roc_auc_score(y_truth, y_pred)
        return auc_score

    def test_auc(self, datas, model, saveData=False):
        res, tar = self.get_scores_from_model(datas, model)
        a = self.auc(tar, res)
        if saveData:
            file = shelve.open(
                r"/home/cyw/projects/function_sim_project/all_data/indicators/auc/siamese_img_auc")
            file["pred"] = res
            file["actual"] = tar
            file.close()
        return a

    def get_scores_from_model(self, datas, model):
        model.eval()
        with torch.no_grad():
            device = torch.device(
                "cuda:0" if torch.cuda.is_available() else "cpu")
            model.to(device)
            res = []
            tar = []
            lth = len(datas)
            for i in range(lth):
                x, y, sim = datas[i][0], datas[i][1], datas[i][2]
                x = x.to(device, torch.float)
                y = y.to(device, torch.float)
                sim = sim.to(device, torch.float)
                pred = model(x, y)
                # res.extend([1 if i>0.5 else 0 for i in list(pred)])
                res.extend(pred.cpu().numpy().tolist())
                tar.extend([1 if i > 0.5 else 0 for i in list(sim)])
        return res, tar


if __name__ == "__main__":
    # nohup python3 /home/cyw/projects/function_sim_project/Siamese_with_img/siamese_img_main.py
    trainer = siamese_img_train()
    trainer.train()
